#include <cuda_runtime.h>
#include <cublas_v2.h>
#include <stdio.h>

int main() {
    cublasHandle_t handle;
    cublasCreate(&handle);

    // 定义矩阵 A 和 B
    // A: 2×3 (行主序)
    float A[6] = {1, 2, 3, 4, 5, 6}; 
    // // B: 4×3 (行主序) 
    float B[12] = {7, 8, 9, 10, 11, 12, 13, 14, 15,16,17,18}; 
    // C: 2×4 结果矩阵
    float C[8] = {0};
    float alpha = 1.0f, beta = 0.0f;

    // 分配 GPU 内存
    float *d_A, *d_B, *d_C;
    cudaMalloc(&d_A, sizeof(A));
    cudaMalloc(&d_B, sizeof(B));
    cudaMalloc(&d_C, sizeof(C));

    // 将数据从主机复制到设备
    cudaMemcpy(d_A, A, sizeof(A), cudaMemcpyHostToDevice);
    cudaMemcpy(d_B, B, sizeof(B), cudaMemcpyHostToDevice);
    cudaMemset(d_C, 0, sizeof(C));

    // 调用 cublasSgemm
    // 计算 C = A * B^T
    cublasSgemm(handle, CUBLAS_OP_T, CUBLAS_OP_N, 4, 2, 3, &alpha, d_B, 3, d_A, 3, &beta, d_C, 4);


    // 将结果从设备复制回主机
    cudaMemcpy(C, d_C, sizeof(C), cudaMemcpyDeviceToHost);

    // 输出结果矩阵 C
    printf("Result matrix C:\n");
    for (int i = 0; i < 2; ++i) {
        for (int j = 0; j < 4; ++j) {
            printf("%f ", C[i * 4 + j]);
        }
        printf("\n");
    }

    // 释放 GPU 内存
    cudaFree(d_A);
    cudaFree(d_B);
    cudaFree(d_C);

    // 销毁 cuBLAS 上下文
    cublasDestroy(handle);

    return 0;
}